import pyatspi
import logging
from lxml import etree  # 用于 XPath 查询和 XML 操作
from typing import Dict, Optional, Tuple, List, Set, TypeAlias
from types import SimpleNamespace
from pprint import pprint
from utils import setup_logger

class JavaGUIManager:

    def __init__(self, title: str, desktop_index: int = 0, logger: logging.Logger = None):
        self.title = title
        self.desktop_index = desktop_index
        self.logger = logger or setup_logger('TA')
        self.logger.info(f'查找 标题包含 "{self.title}" 的 java 程序')
        self.logger.info('请确认启动java程序的时候使用了 -Djavax.accessibility.assistive_technologies=org.GNOME.Accessibility.AtkWrapper 选项')
        self.logger.info('请确认启动java程序的时候使用了 -cp /usr/share/java/java-atk-wrapper.jar 参数')
        self.component_cache: Dict[str, pyatspi.Accessible] = {}  # type: ignore # 缓存组件对象
        self.name_index: Dict[str, Set[str]] = {}  # 按 name 构建索引
        self.role_index: Dict[str, Set[str]] = {}  # 按 role 构建索引
        self.buttons = []
        self.text_boxs = []
        self.all_compoments = []
        self.app = None

    def scan(self):
        self.app = self.find_app_by_name(app_name=self.title, desktop_index=self.desktop_index)
        self.scan_components(self.app)
        self.show_summary()

    def get_desktop_count(self) -> int:
        """
        获取当前桌面的数量。
        """
        return pyatspi.Registry.getDesktopCount()

    def find_app_by_name(self, app_name: str, desktop_index: int = 0) -> Optional[pyatspi.Accessible]:  # type: ignore
        """
        根据应用程序名称和桌面编号查找匹配的应用程序。
        """
        desktop_count: int = self.get_desktop_count()
        self.logger.info(f'共有 {desktop_count} 个桌面')
        if desktop_index >= desktop_count:
            self.logger.error(f"错误：桌面编号 {desktop_index} 超出范围。当前共有 {desktop_count} 个桌面。")
            return None
        
        desktop: pyatspi.Accessible = pyatspi.Registry.getDesktop(desktop_index)  # type: ignore
        for app in desktop:
            if app.name.strip():
                self.logger.info(f'发现窗口： "{app.name}"')
        for app in desktop:
            if app_name.lower() in app.name.lower():
                self.logger.info(f'找到了标题包含 "{self.title}" 的 APP')
                return app
        raise Exception(f'无法找到标题包含 "{self.title}" 的 APP')

    def store_component(self, component):
        def _get_size():
            p = component.get_position(pyatspi.XY_SCREEN)
            item.position = p.x, p.y
            q = component.get_size()
            item.size = q.x, q.y
            item.center = p.x + q.x // 2, p.y + q.y // 2
            
        # 生成唯一标识符（基于层级路径）
        role: str = component.get_role_name().lower().replace(" ", "_")  # 角色转换为小写并替换空格

        item = SimpleNamespace(obj=component, name=component.name, path=component.path)
        self.all_compoments.append(SimpleNamespace(obj=component, path=component.path, name=component.name, role=role))
        self.logger.debug(f'path: {component.path}, {component}')

        if component.get_role_name() == 'push button':
            item.index = len([x for x in self.buttons if x.name == component.name]) + 1
            _get_size()
            self.buttons.append(item)

        if component.get_role_name() == 'text':
            item.index = len([x for x in self.buttons if x.name == component.name]) + 1
            _get_size()
            item.text = self._get_textbox_content(component)
            self.text_boxs.append(item)
        # self.logger.info(item)

    def show_summary(self):
        self.logger.info('控件扫描总结如下：')
        self.logger.info(f'     一共有 {len(self.buttons)} 个 按钮')
        for button in self.buttons:
            self.logger.debug(f'        按钮 path={button.path}, Name="{button.name}", index={button.index}, '
                              f'位置: {button.position}, 长宽：{button.size}, '
                              f'中心点：{button.center}')
            
        self.logger.info(f'     一共有 {len(self.text_boxs)} 个 输入框')
        for text_box in self.text_boxs:
            self.logger.debug(f'        输入框 , path={text_box.path}, Name="{text_box.name}", index={text_box.index}, 内容："{text_box.text}", '
                              f'位置: {text_box.position}, 长宽：{text_box.size}, '
                              f'中心点：{text_box.center}')

    def scan_components(self, component: pyatspi.Accessible):  # type: ignore
        """
        递归扫描组件，生成唯一标识符并缓存组件对象，同时构建索引。
        """
        self.store_component(component=component)

        # 递归处理子组件
        try:
            for i, child in enumerate(component):
                self.scan_components(child)
        except Exception as e:
            self.loggers.error(f"无法处理组件 {component.name} 的子组件: {e}")

    def _get_textbox_content(self, text_box):
        interface = text_box.queryText()
        return interface.getText(0, interface.characterCount)

    def _get_compoment_by_path(self, path: str):
        compoment = [x for x in self.all_compoments if x.path == path]
        if compoment:
            return compoment[0].obj
        return None

    def click_button(self, nameorpath: str, index: int = 1):
        compoment = self._get_compoment_by_path(nameorpath)
        if compoment is None:
            compoment = [x for x in self.buttons if x.name == nameorpath and x.index == index]
            if compoment:
                compoment = compoment[0].obj
            else:
                compoment = None

        if compoment is None:
            raise Exception(f'无法找到 name={nameorpath} index={index}的按钮')
        else:
            self.logger.info(f'执行操作 ==> 点击按钮 nameorpath="{nameorpath}" index={index}')
            compoment.do_action(0)
            return compoment

    def set_text_box(self, new_value, nameorpath, index: int = 1):
        compoment = self._get_compoment_by_path(nameorpath)
        if compoment is None:
            compoment = [x for x in self.text_boxs if x.name == nameorpath and x.index == index]
            if compoment:
                compoment = compoment[0].obj
            else:
                compoment = None

        if compoment is None:
            raise Exception(f'无法找到 name={nameorpath} index={index}的输入框')
        else:
            curr_value = self._get_textbox_content(compoment)
            self.logger.info(f'执行操作 ==> 设置 name="{nameorpath}" index={index} 当前内容为: "{curr_value}" 的输入框的内容为：{new_value}')
            compoment.set_text_contents(new_value)
            return compoment

    def get_text_box_content(self, name, index: int = 1):
        compoment = [x for x in self.text_boxs if x.name == name and x.index == index]
        if compoment:
            self.logger.info(f'执行操作 ==> 读取 name="{name}" index={index} 的输入框的内容')
            return self._get_textbox_content(compoment[0].obj)
        else:
            raise Exception(f'无法找到 name={name} index={index}的输入框')

    def is_compoment_visible(self, compoment):
        """
        检查控件是否可见。
        """
        state = compoment.getState()
        return pyatspi.STATE_VISIBLE in state.getStates()

    def is_ecompoment_enabled(self, compoment):
        """
        检查控件是否可操作。
        """
        state = compoment.getState()
        return pyatspi.STATE_ENABLED in state.getStates()


if __name__ == "__main__":
    import argparse

    # 解析命令行参数
    parser: argparse.ArgumentParser = argparse.ArgumentParser(description="扫描应用程序的 GUI 并缓存组件。")
    parser.add_argument("app_name", type=str, help="应用程序名称（或部分名称）")
    parser.add_argument("--desktop", type=int, default=0, help="桌面编号（从 0 开始），默认为 0")
    args: argparse.Namespace = parser.parse_args()

    gui = JavaGUIManager(args.app_name)
    gui.scan()
    item = gui.set_text_box('python', '', 1)
    print(gui.get_text_box_content('', 1))
    # print(item)
    # print(dir(item.queryText()))
    # print(dir(item))
    # gui.click_button('Submit')
    # gui.click_button('/org/a11y/atspi/accessible/9')

    # 调用主函数
    # main(args.app_name, args.desktop)